import { Camera, Texture2D } from "cc";
import { ISensor, ObservationType } from "./ISensor";
import { BuiltInSensorType, IBuiltInSensor } from "./IBuiltInSensor";
import { IDisposable } from "../Utils/IDisposable";
import { ObservationSpec } from "./ObservationSpec";
import { CompressionSpec, SensorCompressionType } from "./CompressionSpec";
import { Texture2DWrapper } from "../Utils/Texture2DWrapper";
import { ObservationWriter, ObservationWriterExtension } from "./ObservationWriter";
import { Utilities } from "../Utilities";

/**
 * @en
 *  A sensor that wraps a Camera object to generate visual observations for an agent.
 * 
 * @zh
 *  一个传感器，用于包装一个相机对象，为代理生成视觉观察。
 * 
 */
export class CameraSensor implements ISensor, IBuiltInSensor, IDisposable {
    private _camera: Camera;
    private _width: number;
    private _height: number;
    private _grayscale: boolean;
    private _name: string;
    private _observationSpec: ObservationSpec;
    private _compressionType: SensorCompressionType;
    private _texture: Texture2D;

    public get camera(): Camera {
        return this._camera;
    }

    public set camera(value: Camera) {
        this._camera = value;
    }

    public get compressionType(): SensorCompressionType {
        return this._compressionType;
    }

    public set compressionType(value: SensorCompressionType) {
        this._compressionType = value;
    }

    constructor(camera: Camera, width: number, height: number, grayscale: boolean, name: string, compression: SensorCompressionType, observationType: ObservationType = ObservationType.Default) {
        this._camera = camera;
        this._width = width;
        this._height = height;
        this._grayscale = grayscale;
        this._name = name;
        const channels = grayscale ? 1 : 3;
        this._observationSpec = ObservationSpec.visual(height, width, channels, observationType);
        this._compressionType = compression;


        this._texture = new Texture2D();
        this._texture.reset({
            width: width,
            height: height,
            format: Texture2D.PixelFormat.RGBA8888,
            mipmapLevel: 1,
            baseLevel: 0,
            maxLevel: 100
        })
    }

    public getName(): string {
        return this._name;
    }

    public getObservationSpec(): ObservationSpec {
        return this._observationSpec;
    }

    public getCompressedObservation(): Uint8Array {
        CameraSensor.observationToTexture(this._camera, this._texture, this._width, this._height);
        const compressed = Texture2DWrapper.encodeToPNG(this._texture) ?? new Uint8Array();
        return compressed;
    }

    public write(writer: ObservationWriter): number {

        CameraSensor.observationToTexture(this._camera, this._texture, this._width, this._height);
        const numWritten = ObservationWriterExtension.writeTexture(writer, this._texture, this._grayscale);
        return numWritten;
    }

    public update(): void { }

    public reset(): void { }

    public getCompressionSpec(): CompressionSpec {
        return new CompressionSpec(this._compressionType);
    }

    /**
     * @en
     * Renders a Camera instance to a 2D texture at the corresponding resolution.
     * 
     * @zh
     *  
     * 
     * @param obsCamera 
     * @param texture2D 
     * @param width 
     * @param height 
     */
    public static observationToTexture(obsCamera: Camera, texture2D: Texture2D, width: number, height: number): void {

        const texture = obsCamera.targetTexture;
        if (texture) {
            const pixels = texture.readPixels(0, 0, width, height);
            if (pixels) {
                texture2D.uploadData(pixels);
            }
        }

    }

    public getBuiltInSensorType(): BuiltInSensorType {
        return BuiltInSensorType.CameraSensor;
    }

    public dispose(): void {
        if (this._texture != null) {
            Utilities.destroyTexture(this._texture);
            this._texture = null!;
        }
    }
}